#pragma once
#include "Ray.h"
#include "Scene.h"
#include "Heightmap.h"
#include <stdio.h>

template <class CellType>
class QuadTreeCellSceneNode {
private:
    QuadTreeCellSceneNode<CellType>* children[4];
    CellType* leaf;
    float minHeight, maxHeight;
    int minX, maxX, minY, maxY;

public:
    //Traces a ray into the scene. Returns the intersection.
    Intersection traceRay(Ray& ray) const;

    //Traces a shadow ray. Returns true iff there is a collision.
    bool traceShadowRay(Ray& ray) const;
    
    static QuadTreeCellSceneNode<CellType>* createFromHeightmap(Heightmap& heightmap);
    static QuadTreeCellSceneNode<CellType>* createFromHeightmap(Heightmap& heightmap, int minX, int maxX, int minY, int maxY);
};

template <class CellType>
Intersection QuadTreeCellSceneNode<CellType>::traceRay(Ray& ray) const {
    if (leaf) {
        return leaf->traceRay(ray);
    }
    Ray tracedRay = ray;

    //hits this node?
    float minTx = (minX - ray.origin.x) / ray.dir.x;
    float maxTx = (maxX - ray.origin.x) / ray.dir.x;
    if (minTx > maxTx) swap(minTx, maxTx);

    float minTy = (minY - ray.origin.y) / ray.dir.y;
    float maxTy = (maxY - ray.origin.y) / ray.dir.y;
    if (minTy > maxTy) swap(minTy, maxTy);

    if (minTx > maxTy || minTy > maxTx) return Intersection();

    if (ray.dir.z < 0) {
        float minTz = (maxHeight - ray.origin.z) / ray.dir.z;
        if (minTz > maxTx || minTz > maxTy) return Intersection();
    } else if (ray.dir.z > 0) {
        float maxTz = (maxHeight - ray.origin.z) / ray.dir.z;
        if (minTx > maxTz || minTy > maxTz) return Intersection();
    } else {
        if (ray.origin.z > maxHeight) return Intersection();
    }

    Intersection result;
    for (int i = 0; i < 4; i++) {
        if (children[i]) {
            Intersection hit = children[i]->traceRay(tracedRay);
            if (hit.exists()) {
                result = hit;
                tracedRay.maxT = hit.t;
            }
        }
    }
    return result;
}

template <class CellType>
bool QuadTreeCellSceneNode<CellType>::traceShadowRay(Ray& ray) const {
    return traceRay(ray).exists();
}

template <class CellType>
QuadTreeCellSceneNode<CellType>* QuadTreeCellSceneNode<CellType>::createFromHeightmap(Heightmap& heightmap) {
    return QuadTreeCellSceneNode<CellType>::createFromHeightmap(heightmap, 0, heightmap.width() - 1, 0, heightmap.height() - 1);
}

template <class CellType>
QuadTreeCellSceneNode<CellType>* QuadTreeCellSceneNode<CellType>::createFromHeightmap(Heightmap& heightmap, int minX, int maxX, int minY, int maxY) {
    QuadTreeCellSceneNode<CellType>* result = new QuadTreeCellSceneNode<CellType>();
    result->minX = minX;
    result->maxX = maxX;
    result->minY = minY;
    result->maxY = maxY;
    result->children[0] = 0;
    result->children[1] = 0;
    result->children[2] = 0;
    result->children[3] = 0;
    result->leaf = 0;

    if (maxX > minX + 1) {
        if (maxY > minY + 1) {
            //4-node
            int midX = (minX + maxX) / 2;
            int midY = (minY + maxY) / 2;
            result->children[0] = createFromHeightmap(heightmap, minX, midX, minY, midY);
            result->children[1] = createFromHeightmap(heightmap, midX, maxX, minY, midY);
            result->children[2] = createFromHeightmap(heightmap, minX, midX, midY, maxY);
            result->children[3] = createFromHeightmap(heightmap, midX, maxX, midY, maxY);
            result->minHeight = min(result->children[0]->minHeight, result->children[1]->minHeight);
            result->minHeight = min(result->minHeight, result->children[2]->minHeight);
            result->minHeight = min(result->minHeight, result->children[3]->minHeight);
            result->maxHeight = max(result->children[0]->maxHeight, result->children[1]->maxHeight);
            result->maxHeight = max(result->maxHeight, result->children[2]->maxHeight);
            result->maxHeight = max(result->maxHeight, result->children[3]->maxHeight);
        } else {
            //2-node, split on x
            int midX = (minX + maxX) / 2;
            result->children[0] = createFromHeightmap(heightmap, minX, midX, minY, maxY);
            result->children[1] = createFromHeightmap(heightmap, midX, maxX, minY, maxY);
            result->minHeight = min(result->children[0]->minHeight, result->children[1]->minHeight);
            result->maxHeight = max(result->children[0]->maxHeight, result->children[1]->maxHeight);
        }
    } else {
        if (maxY > minY + 1) {
            //2-node, split on y
            int midY = (minY + maxY) / 2;
            result->children[0] = createFromHeightmap(heightmap, minX, maxX, minY, midY);
            result->children[3] = createFromHeightmap(heightmap, minX, maxX, midY, maxY);
            result->minHeight = min(result->children[0]->minHeight, result->children[3]->minHeight);
            result->maxHeight = max(result->children[0]->maxHeight, result->children[3]->maxHeight);
        } else {
            //leaf node
            result->leaf = new CellType();
            result->leaf->setFromHeightmap(heightmap, minX, minY);
            result->minHeight = result->leaf->getMinHeight();
            result->maxHeight = result->leaf->getMaxHeight();
        }
    }
    return result;
}


template <class CellType>
class QuadTreeCellScene :
    public Scene {
private:
    CellType* cells;

    QuadTreeCellSceneNode<CellType>* root;
public:
    //Readies this scene for ray tracing.
    void ready();

    //Assigns a heightmap to this scene.
    void assignHeightmap(Heightmap& heightmap);

    //Traces a ray into the scene. Returns the intersection.
    Intersection traceRay(Ray& ray) const;

    //Traces a shadow ray. Returns true iff there is a collision.
    bool traceShadowRay(Ray& ray) const;
};

template <class CellType>
void QuadTreeCellScene<CellType>::ready() {
    
}

template <class CellType>
void QuadTreeCellScene<CellType>::assignHeightmap(Heightmap& heightmap) {
    //construct quadtree
    root = QuadTreeCellSceneNode<CellType>::createFromHeightmap(heightmap); 
}

template <class CellType>
Intersection QuadTreeCellScene<CellType>::traceRay(Ray& ray) const {
    return root->traceRay(ray);
}

template <class CellType>
bool QuadTreeCellScene<CellType>::traceShadowRay(Ray& ray) const {
    return root->traceShadowRay(ray);
}